import torch
import torch.nn as nn
from torch.nn import functional as F
import numpy as np
from torch.nn.parameter import Parameter

'''
用于计算一批点对之间的对极几何误差（epipolar symmetry error）。
这个函数通常用于计算机视觉中的立体匹配或运动恢复结构（Structure from Motion）任务中，以评估点对在两个视图中是否匹配。
以下是对代码的详细解释：
这两个参数是形状为 (batch_size, num_pts, 2) 的张量，表示在两个视图中的点坐标
F:这是一个形状为 (batch_size, 3, 3) 的张量，表示基本矩阵（Fundamental Matrix），用于描述两个视图之间的对极几何关系.
'''
def batch_episym(x1, x2, F):
    # 扩展点坐标：
    batch_size, num_pts = x1.shape[0], x1.shape[1]
    #  将 x1 和 x2 的每个点坐标扩展为齐次坐标形式，即在最后添加一个1，然后重置为张量模式
    x1 = torch.cat([x1, x1.new_ones(batch_size, num_pts,1)], dim=-1).reshape(batch_size, num_pts,3,1)
    x2 = torch.cat([x2, x2.new_ones(batch_size, num_pts,1)], dim=-1).reshape(batch_size, num_pts,3,1)
    # 扩展基本矩阵：
    # 将基本矩阵 F 扩展为 (batch_size, num_pts, 3, 3) 的形状，以便与每个点的坐标进行运算.
    # repeat：某个维度重复多次
    F = F.reshape(-1,1,3,3).repeat(1,num_pts,1,1)

    x2Fx1 = torch.matmul(x2.transpose(2,3), torch.matmul(F, x1)).reshape(batch_size,num_pts)
    Fx1 = torch.matmul(F,x1).reshape(batch_size,num_pts,3)
    Ftx2 = torch.matmul(F.transpose(2,3),x2).reshape(batch_size,num_pts,3)

    ys = x2Fx1**2 * (
            1.0 / (Fx1[:, :, 0]**2 + Fx1[:, :, 1]**2 + 1e-15) +
            1.0 / (Ftx2[:, :, 0]**2 + Ftx2[:, :, 1]**2 + 1e-15))

    return ys

def knn(x, k):
    inner = -2*torch.matmul(x.transpose(2, 1), x)
    xx = torch.sum(x**2, dim=1, keepdim=True)
    pairwise_distance = -xx - inner - xx.transpose(2, 1)

    idx = pairwise_distance.topk(k=k, dim=-1)[1]   # (batch_size, num_points, k)

    return idx[:, :, :]

def get_graph_feature(x, k=20, idx=None):
    batch_size = x.size(0)
    num_points = x.size(2)
    x = x.view(batch_size, -1, num_points)
    if idx is None:
        idx_out = knn(x, k=k)
    else:
        idx_out = idx
    device = x.device

    idx_base = torch.arange(0, batch_size, device=device).view(-1, 1, 1)*num_points

    idx = idx_out + idx_base

    idx = idx.view(-1)

    _, num_dims, _ = x.size()

    x = x.transpose(2, 1).contiguous()
    feature = x.view(batch_size*num_points, -1)[idx, :]
    feature = feature.view(batch_size, num_points, k, num_dims)
    x = x.view(batch_size, num_points, 1, num_dims).repeat(1, 1, k, 1)
    feature = torch.cat((x, x - feature), dim=3).permute(0, 3, 1, 2).contiguous()
    return feature

'''
定义了一个残差块（Residual Block），是残差网络（ResNet）中的基本构建模块。
残差块通过引入“跳跃连接”（skip connection）来缓解深层神经网络中的梯度消失问题，从而使得网络可以更深。以下是对该类的详细解释：
'''
class ResNet_Block(nn.Module):
    '''
    pre用于决定是否使用 self.right 路径进行前处理.
    '''
    def __init__(self, inchannel, outchannel, pre=False):
        super(ResNet_Block, self).__init__()
        self.pre = pre
        # 定义一个简单的卷积路径，只包含一个 1x1 卷积层，用于调整输入特征图的通道数.
        self.right = nn.Sequential(
            nn.Conv2d(inchannel, outchannel, (1, 1)),
        )
        # 定义一个更复杂的卷积路径，包含两个 1x1 卷积层和相应的归一化层与激活函数，用于提取特征.
        # 第一个卷积层后接 nn.InstanceNorm2d 和 nn.BatchNorm2d，分别进行实例归一化和批量归一化，然后是 ReLU 激活函数.
        # 第二个卷积层后接 nn.InstanceNorm2d 和 nn.BatchNorm2d，但没有激活函数.
        self.left = nn.Sequential(
            nn.Conv2d(inchannel, outchannel, (1, 1)),
            nn.InstanceNorm2d(outchannel),
            nn.BatchNorm2d(outchannel),
            nn.ReLU(),
            nn.Conv2d(outchannel, outchannel, (1, 1)),
            nn.InstanceNorm2d(outchannel),
            nn.BatchNorm2d(outchannel),
        )

    # x 输入特征图.
    def forward(self, x):
        x1 = self.right(x) if self.pre is True else x
        out = self.left(x)
        out = out + x1
        return torch.relu(out)

'''
一个用于估计本质矩阵（Essential Matrix）的函数，它结合了加权的8点算法。以下是对该函数的详细解释：
x_in: 形状为 (batch, 1, N, 4) 的张量，表示输入的匹配点对。其中 N 是匹配点对的数量，每个点对包含4个坐标值（通常是归一化的坐标）.
logits: 形状为 (batch, 2, N, 1) 的张量，包含用于加权的 logits 值。第一个通道用于计算掩码（mask），第二个通道用于计算权重（weights）.
'''
def weighted_8points(x_in, logits):
    # 计算掩码和权重：
    # x_in: batch * 1 * N * 4
    mask = logits[:, 0, :, 0]
    weights = logits[:, 1, :, 0]

    # 将掩码`logits`转换为概率值，范围在`[0, 1]`之间.
    mask = torch.sigmoid(mask)
    # mask：将权重logits转换为权重值，并乘以掩码概率值，得到最终的权重.
    weights = torch.exp(weights) * mask
    # 对权重进行归一化，确保每个样本的权重和为1，防止除以零的情况.
    weights = weights / (torch.sum(weights, dim=-1, keepdim=True) + 1e-5)

    # 数据预处理
    x_shp = x_in.shape
    x_in = x_in.squeeze(1)

    # 将输入张量重新排列为(batch, 4, N) 的形状，以便后续操作.
    # permute换顺序
    xx = torch.reshape(x_in, (x_shp[0], x_shp[2], 4)).permute(0, 2, 1).contiguous()

    # 构建矩阵X：
    # X 的形状为 (batch, N, 9)，表示每个匹配点对对应的线性方程.
    X = torch.stack([
        xx[:, 2] * xx[:, 0], xx[:, 2] * xx[:, 1], xx[:, 2],
        xx[:, 3] * xx[:, 0], xx[:, 3] * xx[:, 1], xx[:, 3],
        xx[:, 0], xx[:, 1], torch.ones_like(xx[:, 0])
    ], dim=1).permute(0, 2, 1).contiguous()
    # 加权矩阵计算
    wX = torch.reshape(weights, (x_shp[0], x_shp[2], 1)) * X
    XwX = torch.matmul(X.permute(0, 2, 1).contiguous(), wX)

    # 估计本质矩阵：
    # 计算对称矩阵 XwX 的特征值和特征向量.
    # Recover essential matrix from self-adjoing eigen
    _, v = torch.linalg.eigh(XwX)

    # 取特征向量的第一个特征向量，作为估计的本质矩阵.
    e_hat = torch.reshape(v[:, :, 0], (x_shp[0], 9))

    # Make unit norm just in case
    # 将估计的本质矩阵归一化，确保其范数为1.
    e_hat = e_hat / torch.norm(e_hat, dim=1, keepdim=True)
    return e_hat

# 定义了一个基于动态图卷积神经网络（Dynamic Graph Convolutional Neural Network, DGCNN）的块。
# DGCNN 是一种用于处理点云数据的深度学习架构，通过构建动态图来捕捉点云中的局部几何结构。以下是对该类的详细解释：
class DGCNN_Block(nn.Module):
    '''
    knn_num：KNN（K-Nearest Neighbors）的邻居数量，默认为9。表示每个点在构建图时考虑的最近邻居数量.
    in_channel：输入通道数，默认为128。表示输入特征的通道数.
    '''
    def __init__(self, knn_num=9, in_channel=128):
        super(DGCNN_Block, self).__init__()
        self.knn_num = knn_num
        self.in_channel = in_channel

        assert self.knn_num == 9 or self.knn_num == 6
        if self.knn_num == 9:
            # 定义一个包含两个卷积层的序列，每个卷积层后面都接有批量归一化层和ReLU激活函数。
            # 第一个卷积层的卷积核大小为 (1, 3)，步长为 (1, 3)，第二个卷积层的卷积核大小为 (1, 3).
            self.conv = nn.Sequential(
                nn.Conv2d(self.in_channel*2, self.in_channel, (1, 3), stride=(1, 3)),
                nn.BatchNorm2d(self.in_channel),
                nn.ReLU(inplace=True),
                nn.Conv2d(self.in_channel, self.in_channel, (1, 3)),
                nn.BatchNorm2d(self.in_channel),
                nn.ReLU(inplace=True),
            )
        if self.knn_num == 6:
            # 则定义一个包含两个卷积层的序列，第一个卷积层与 knn_num == 9 的情况相同，第二个卷积层的卷积核大小为 (1, 2).
            self.conv = nn.Sequential(
                nn.Conv2d(self.in_channel*2, self.in_channel, (1, 3), stride=(1, 3)),
                nn.BatchNorm2d(self.in_channel),
                nn.ReLU(inplace=True),
                nn.Conv2d(self.in_channel, self.in_channel, (1, 2)),
                nn.BatchNorm2d(self.in_channel),
                nn.ReLU(inplace=True),
            )

    # features，输入特征张量，形状为 (B, C, N, K)，其中 B 是批量大小，C 是通道数，N 是点的数量，K 是邻居数量.
    def forward(self, features):
        B, _, N, _ = features.shape
        # 根据输入特征和KNN邻居数量构建图特征。这个函数通常会计算每个点与其最近邻居之间的特征差值，以捕捉局部几何结构.
        out = get_graph_feature(features, k=self.knn_num)
        out = self.conv(out)
        return out

class GCN_Block(nn.Module):
    def __init__(self, in_channel):
        super(GCN_Block, self).__init__()
        self.in_channel = in_channel
        self.conv = nn.Sequential(
            nn.Conv2d(self.in_channel, self.in_channel, (1, 1)),
            nn.BatchNorm2d(self.in_channel),
            nn.ReLU(inplace=True),
        )

    def attention(self, w):
        w = torch.relu(torch.tanh(w)).unsqueeze(-1)
        A = torch.bmm(w, w.transpose(1, 2))
        return A

    def graph_aggregation(self, x, w):
        B, _, N, _ = x.size()
        with torch.no_grad():
            A = self.attention(w)
            I = torch.eye(N).unsqueeze(0).to(x.device).detach()
            A = A + I
            D_out = torch.sum(A, dim=-1)
            D = (1 / D_out) ** 0.5
            D = torch.diag_embed(D)
            L = torch.bmm(D, A)
            L = torch.bmm(L, D)
        out = x.squeeze(-1).transpose(1, 2).contiguous()
        out = torch.bmm(L, out).unsqueeze(-1)
        out = out.transpose(1, 2).contiguous()

        return out

    def forward(self, x, w):
        out = self.graph_aggregation(x, w)
        out = self.conv(out)
        return out

'''
动态块
'''
class DS_Block(nn.Module):
    def __init__(self, initial=False, predict=False, out_channel=128, k_num=8, sampling_rate=0.5):
        super(DS_Block, self).__init__()
        self.initial = initial
        self.in_channel = 4 if self.initial is True else 6
        self.out_channel = out_channel
        self.k_num = k_num
        self.predict = predict
        self.sr = sampling_rate

        self.conv = nn.Sequential(
            nn.Conv2d(self.in_channel, self.out_channel, (1, 1)),
            # 批量归一化（正则化）的作用是减少内部协变量偏移，加速训练过程，并有助于提高模型的泛化能力.
            nn.BatchNorm2d(self.out_channel),
            # inplace 表示在原地进行操作，即直接在输入张量上进行修改，而不是创建一个新的张量。这样可以节省内存.
            nn.ReLU(inplace=True)
        )

        # 全局连接学习
        self.gcn = GCN_Block(self.out_channel)

        # 一堆残差网络的处理（不用做简单处理）
        # 动态图卷积神经网络
        self.embed_0 = nn.Sequential(
            ResNet_Block(self.out_channel, self.out_channel, pre=False),
            ResNet_Block(self.out_channel, self.out_channel, pre=False),
            ResNet_Block(self.out_channel, self.out_channel, pre=False),
            ResNet_Block(self.out_channel, self.out_channel, pre=False),
            # 动态图卷积神经网络
            DGCNN_Block(self.k_num, self.out_channel),
            ResNet_Block(self.out_channel, self.out_channel, pre=False),
            ResNet_Block(self.out_channel, self.out_channel, pre=False),
            ResNet_Block(self.out_channel, self.out_channel, pre=False),
            ResNet_Block(self.out_channel, self.out_channel, pre=False),
        )
        self.embed_1 = nn.Sequential(
            ResNet_Block(self.out_channel, self.out_channel, pre=False),
        )
        self.linear_0 = nn.Conv2d(self.out_channel, 1, (1, 1))
        self.linear_1 = nn.Conv2d(self.out_channel, 1, (1, 1))

        if self.predict == True:
            self.embed_2 = ResNet_Block(self.out_channel, self.out_channel, pre=False)
            self.linear_2 = nn.Conv2d(self.out_channel, 2, (1, 1))

    def down_sampling(self, x, y, weights, indices, features=None, predict=False):
        B, _, N , _ = x.size()
        # 只使用权重很强的几个位置
        indices = indices[:, :int(N*self.sr)]
        with torch.no_grad():
            # 从输入张量中根据索引选择元素。它在指定的维度上进行操作，将索引对应的元素收集到一个新的张量中。
            # 结果张量的形状与 index 的形状相同.
            y_out = torch.gather(y, dim=-1, index=indices)
            w_out = torch.gather(weights, dim=-1, index=indices)
        # 改变张量的形状
        indices = indices.view(B, 1, -1, 1)

        if predict == False:
            with torch.no_grad():
                x_out = torch.gather(x[:, :, :, :4], dim=2, index=indices.repeat(1, 1, 1, 4))
            return x_out, y_out, w_out
        else:
            with torch.no_grad():
                x_out = torch.gather(x[:, :, :, :4], dim=2, index=indices.repeat(1, 1, 1, 4))
            feature_out = torch.gather(features, dim=2, index=indices.repeat(1, 128, 1, 1))
            return x_out, y_out, w_out, feature_out

    def forward(self, x, y):
        # 处理x
        # batch Chanel，X, y
        B, _, N , _ = x.size()
        out = x.transpose(1, 3).contiguous()
        out = self.conv(out)

        out = self.embed_0(out)
        w0 = self.linear_0(out).view(B, -1)

        out_g = self.gcn(out, w0.detach())
        out = out_g + out

        out = self.embed_1(out)
        w1 = self.linear_1(out).view(B, -1)

        # 提取结果（y登场）
        if self.predict == False:
            # 权重排序：降序排序，返回降序的结果以及原函数的索引
            w1_ds, indices = torch.sort(w1, dim=-1, descending=True)
            w1_ds = w1_ds[:, :int(N*self.sr)]
            # TAG 这里的w0是之前的一次运算后的结果，我们这边要根据后边做的实验获取的优先的点位indices获取前一次获取的权重的结果
            # indices就是后一次实验获得的索引值
            x_ds, y_ds, w0_ds = self.down_sampling(x, y, w0, indices, None, self.predict)

            return x_ds, y_ds, [w0, w1], [w0_ds, w1_ds]
        else:
            w1_ds, indices = torch.sort(w1, dim=-1, descending=True)
            w1_ds = w1_ds[:, :int(N*self.sr)]
            # TODO 这里的w0是之前的一次运算后的结果，我们这边要将后边做的实验的与这次的实验做一个合并运算，indices就是后一次实验获得的索引值
            # 其实就是说第二次做的实验后我获得了一些优先的值，然后我想根据前一个实验做的值再做一次优化
            # 啊，呸！就是获取前一次实验的权重结果
            # #################################
            # TAG 重点！！！！完结了，
            # 让我缓一缓。。。。。。
            # 欧克
            # out一个临时存储值的地方 ：这里的是指最新的训练结果，提取高优先的out给out
            x_ds, y_ds, w0_ds, out = self.down_sampling(x, y, w0, indices, out, self.predict)
            # 做训练
            out = self.embed_2(out)
            w2 = self.linear_2(out)
            # 加权8点运算，去考虑临近值
            e_hat = weighted_8points(x_ds, w2)

            # xy的高权点位，[0，1，2训练权值],[0，1的高权值，（1为顺序排列）]，相关值
            return x_ds, y_ds, [w0, w1, w2[:, 0, :, 0]], [w0_ds, w1_ds], e_hat

# NOTE CLNet
class CLNet(nn.Module):
    def __init__(self, config):
        super(CLNet, self).__init__()

        self.ds_0 = DS_Block(initial=True, predict=False, out_channel=128, k_num=9, sampling_rate=config.sr)#1.0)
        self.ds_1 = DS_Block(initial=False, predict=True, out_channel=128, k_num=6, sampling_rate=config.sr)

    def forward(self, x, y):
        B, _, N, _ = x.shape

        # 9个临近点的获取
        x1, y1, ws0, w_ds0 = self.ds_0(x, y)

        # 取样权重，B数组的数据重置为（1，N,1）,压缩圈中
        w_ds0[0] = torch.relu(torch.tanh(w_ds0[0])).reshape(B, 1, -1, 1)
        # 训练的权重
        w_ds0[1] = torch.relu(torch.tanh(w_ds0[1])).reshape(B, 1, -1, 1)
        # detach用于从当前的计算图中分离一个张量
        x_ = torch.cat([x1, w_ds0[0].detach(), w_ds0[1].detach()], dim=-1)

        x2, y2, ws1, w_ds1, e_hat = self.ds_1(x_, y1)

        with torch.no_grad():
            # 几何误差计算
            y_hat = batch_episym(x[:, 0, :, :2], x[:, 0, :, 2:], e_hat)
        # 权重，数组 e_hat用于表示残差的估计值，即实际观测值与拟合值之间的差异
        return ws0 + ws1, [y, y, y1, y1, y2], [e_hat], y_hat
